from cloudcontrol._async import AsyncIterable, AsyncIterableGroup
from threading import Lock


class GeneratorStream(object):
    
    def __init__(self, generator, input_handler, kill_handler):
        self.__async_iterable = AsyncIterable( generator )
        self.__input_handler = input_handler
        self.__kill_handler = kill_handler
        self.__input_lock = Lock()
    
    def __iter__(self):
        for obj in self.__async_iterable:
            yield obj
    
    def inject(self, obj):
        # TODO: this may need some kind of special object (None may not be appropriate) to indicate when no more input will be provided?
        self.__input_lock.acquire()
        try:
            self.__input_handler( obj )
        finally:
            self.__input_lock.release()
        return self
    
    def drain(self):
        for dontcare in self:
            pass
    
    def kill(self):
        self.__kill_handler()
        self.__async_iterable.kill()


class CommandStream(object):
    
    def __init__(self, stdout_generator, stderr_generator, input_handler, kill_handler):
        self.__async_iterable = AsyncIterableGroup( stdout_generator, stderr_generator )
        self.__input_handler = input_handler
        self.__kill_handler = kill_handler
        self.__input_lock = Lock()
    
    # FIXME: make a "bytes" iterator
    # FIXME: yield subclasses of tuple with field names
    
    # TODO: remove this method, it sucks
    def get_fake_stdin(self):
        class FakeStdin(object):
            def __init__(self, command_stream):
                self.__command_stream = command_stream
            def write(self, buffer_str):
                self.__command_stream.inject(buffer_str)
        return FakeStdin(self)
    
    def __iter__(self):
        """yields each byte of the stream, and also
        (as a convenience) each line of the stream"""
        pending_stdout_line = ""
        pending_stderr_line = ""
        for stdout_byte, stderr_byte in self.__async_iterable:
            stdout_line = None
            stderr_line = None
            if stdout_byte:
                if stdout_byte == "\n":
                    stdout_line = pending_stdout_line
                    pending_stdout_line = ""
                    if stdout_line.endswith("\r"):
                        stdout_line = stdout_line[0:-1]
                else:
                    pending_stdout_line += stdout_byte
            if stderr_byte:
                if stderr_byte == "\n":
                    stderr_line = pending_stderr_line
                    pending_stderr_line = ""
                    if stderr_line.endswith("\r"):
                        stderr_line = stderr_line[0:-1]
                else:
                    pending_stderr_line += stderr_byte
            
            # FIXME: rather than an N-tuple, use the 
            yield stdout_line, stderr_line, stdout_byte, stderr_byte
    
    def inject(self, buffer_str):
        self.__input_lock.acquire()
        try:
            self.__input_handler( buffer_str )
        finally:
            self.__input_lock.release()
        return self
    
    def drain(self):
        for dontcare in self:
            pass
    
    def kill(self):
        self.__kill_handler()
        self.__async_iterable.kill()
